from logger_util import build_logger
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
import time
import threading
import json

worker = None
logger = None

app = FastAPI()


def heart_beat_worker(obj):
    while True:
        time.sleep(5)
        logger.info(f'thread: Heartbeat from Worker {obj.worker_id}!')
        obj.send_heart_beat()


class BaseModelWorker:
    def __init__(
        self,
        worker_id: str,
    ):
        global logger, worker

        self.worker_id = worker_id

        self.heart_beat_thread = None

        if logger is None:
            # logger = build_logger("model_worker", f"model_worker_{self.worker_id}.log")
            logger = build_logger("base_model_worker", f"base_model_worker_{self.worker_id}.log")
        if worker is None:
            worker = self
            
        self.init_heart_beat()
            
    def init_heart_beat(self):
        self.heart_beat_thread = threading.Thread(
            target=heart_beat_worker,
            args=(self,),
            daemon=True,
        )
        self.heart_beat_thread.start()
        
    def send_heart_beat(self):
        logger.info(f'self.send_heart_beat: Heartbeat from Worker {self.worker_id}!')
        
        
    def generate_stream_gate(self, params):
        for i in range(20):
            logger.info(f"{params} => data: {i}\n\n")
            res_dict = {
                "input": params,
                "output": i,
            }
            yield json.dumps(res_dict, ensure_ascii=False).encode('utf8') + b"\0"
            time.sleep(0.2)
        
@app.post("/worker_generate_stream")
async def api_generate_stream(request: Request):
    params = await request.json()
    generator = worker.generate_stream_gate(params)
    return StreamingResponse(generator)